'''
Created on 2009-7-19

@author: madlee
'''
import os.path
import sqlite3
from datetime import datetime

from kuai import ForceField, EquivalenceTable
from kuai.ff import ForceFieldKey, ForceFieldItem, EquivalenceTableItem


DFFDB_SQL_CREATE_TABLE = ["""CREATE TABLE IF NOT EXISTS 'DBINDEX'
(id INTEGER PRIMARY KEY AUTOINCREMENT, 
name TEXT UNIQUE, type TEXT, 
birth DATETIME DEFAULT CURRENT_TIMESTAMP,
lmt DATETIME DEFAULT CURRENT_TIMESTAMP,
license TEXT DEFAULT '');""",

"""INSERT INTO 'DBINDEX' (name, type) VALUES ('%(name)s', '%(type)s')""", 

"""CREATE TABLE '%(name)s' 
(version INTEGER, key TEXT, 
 value TEXT, reference INTEGER DEFAULT NULL, 
 PRIMARY KEY(version, key));"""]

DFFDB_SQL_OPEN_DATABASE = """SELECT id, name, type FROM 'DBINDEX' WHERE name = ?"""
DFFDB_SQL_NEXT_VERSION = """SELECT MAX(version) FROM '%s'"""

DFFDB_SQL_LAST_ITEM = """SELECT version, key, value, reference FROM '%s' WHERE key = ? ORDER BY version DESC"""
DFFDB_SQL_INSERT_ITEM = """INSERT INTO '%s' (version, key, value, reference) VALUES (?, ?, ?, ?)"""

DFFDB_SQL_SELECT_ITEMS = """SELECT MAX(version), key, value, reference from '%s' WHERE version <= %d and key in (%s) GROUP BY key ORDER BY key"""
DFFDB_SQL_TOUCH_TABLE = """UPDATE DBINDEX SET lmt = ? WHERE name = ?"""

DFFDB_SQL_EXPORT_VERSION = """SELECT MAX(version), key, value, reference from '%s' WHERE version <= %d GROUP BY key ORDER BY key""" 

DFFDB_SQL_SELECT_ALL = """SELECT version, key, value, reference from '%s' ORDER BY key, version DESC""" 

class Database:
    def __init__(self, filename, tablename):
        self.filename = filename
        self.tablename = tablename
        self.__conn = sqlite3.connect(filename)
        self.__cursor = self.__conn.cursor()
        self.__cursor.execute(DFFDB_SQL_OPEN_DATABASE,(tablename,))
        result = self.__cursor.fetchone()
        if result is None:
            raise RuntimeError("Table %s is not existed!" % (tablename,))
        else:
            self.id = result[0]
            self.type = result[2]
    
    def next_version(self):
        sql = DFFDB_SQL_NEXT_VERSION % (self.tablename,)
        self.__cursor.execute(sql)
        result = self.__cursor.fetchone()
        if result[0] is None:
            return 1
        else:
            return result[0]+1
        
    def checkin(self, ff, eqt, ref=None):
        if ff.type != self.type:
            raise RuntimeError("The force field type is different to this database")
        newdata = []
        for k, v in eqt.iteritems():
            newvalue = " ".join(v)
            lasteqt = self.__get_last_eqt(k)
            if lasteqt is None or lasteqt[2] != newvalue:
               newdata.append((Database.__format_eqt_key(k), newvalue,))
        for k, v in ff.iteritems():
            newvalue = " ".join([str(i) for i in v])
            lastitem = self.__get_last_item(k)
            if lastitem is None or lastitem[2] != newvalue:
                newdata.append((str(k), newvalue,))

        version = self.next_version()
        if len(newdata) > 0:
            sql = DFFDB_SQL_INSERT_ITEM % (self.tablename) 
            for i in newdata:
                self.__cursor.execute(sql, (version, i[0], i[1], ref))
            self.__cursor.execute(DFFDB_SQL_TOUCH_TABLE, [datetime.utcnow(), self.tablename])
            self.__conn.commit()
            return version, newdata, 
        else:
            return version-1, newdata
    
    def checkout(self, keyiter, version=None):
        if version is None:
            version = self.next_version()
        from sets import Set
        types = Set()
        for i in keyiter:
            for j in i.atoms:
                types.add(Database.__format_eqt_key(j))
        sql = DFFDB_SQL_SELECT_ITEMS % (self.tablename, version, ",".join(['?'] * len(types)))
        print sql
        print [i for i in types]
        self.__cursor.execute(sql, [i for i in types])
        eqt = EquivalenceTable()
        for row in self.__cursor:
            version = row[0]
            key, item = Database.__parse_eqt_row(row)
            eqt[key] = item
        missing_items = Set()
        for i in keyiter:
            missing_items.add(eqt.get_equal_key(i))
        sql = DFFDB_SQL_SELECT_ITEMS % (self.tablename, version, ",".join(['?'] * len(missing_items)))
        self.__cursor.execute(sql, [str(i) for i in missing_items])
        ff = ForceField(self.type)
        for row in self.__cursor:
            key, item = Database.__parse_ppf_row(row)
            ff[key] = item
        return eqt, ff
       
    def export(self, filename, version=None):
        if version is None:
            version = self.next_version()
            
        ppffile = open(filename, "wt")
        eqtfile = open(EquivalenceTable.get_eqt_filename(filename), "wt")
        from kuai.dff.dffppf import PPF_HEADER
        ppffile.write(PPF_HEADER % self.type)
        
        sql = DFFDB_SQL_EXPORT_VERSION % (self.tablename, version)
        self.__cursor.execute(sql)
        for i in self.__cursor:
            if i[1].startswith("EQT:"):
                key, item = Database.__parse_eqt_row(i)
                eqtfile.write(key + ": " + str(item) + "\n")
            else:
                key, item = Database.__parse_ppf_row(i)
                ppffile.write(str(key) + str(item) + "\n")
        
        ppffile.close()
        eqtfile.close()
        
    def dumpall(self, filename):
        ppffile = open(filename, "wt")
        eqtfile = open(EquivalenceTable.get_eqt_filename(filename), "wt")
        from kuai.dff.dffppf import PPF_HEADER
        ppffile.write(PPF_HEADER % self.type)
        
        sql = DFFDB_SQL_SELECT_ALL % (self.tablename)
        self.__cursor.execute(sql)
        for i in self.__cursor:
            if i[1].startswith("EQT:"):
                key, item = Database.__parse_eqt_row(i)
                eqtfile.write(key + ": " + str(item) + "\n")
            else:
                key, item = Database.__parse_ppf_row(i)
                ppffile.write(str(key) + str(item) + "\n")
        
        ppffile.close()
        eqtfile.close()
        
    @staticmethod
    def create(filename, tablename, type):
        conn = sqlite3.connect(filename)
        cursor = conn.cursor()
        for i in DFFDB_SQL_CREATE_TABLE:
            i = i % {'name':tablename, 'type':type}
            cursor.execute(i)
        conn.commit()
        cursor.close()
        conn.close()
        return Database(filename, tablename)
        
        
    def __get_last_eqt(self, key):
        return self.__get_last_item(Database.__format_eqt_key(key))
        
    def __get_last_item(self, key):
        sql = DFFDB_SQL_LAST_ITEM % (self.tablename) 
        self.__cursor.execute(sql, (str(key),))
        return self.__cursor.fetchone()
    
    @staticmethod
    def __format_eqt_key(key):
        return "EQT: " + key
    
    @staticmethod
    def __parse_eqt_row(v):
        assert v[1].startswith("EQT:")
        key = v[1][4:].strip()
        item = EquivalenceTableItem()
        item += v[2].split()
        item.version = v[0]
        item.reference = v[3]
        return key, item
    
    @staticmethod
    def __parse_ppf_row(v):
        tokens = v[1].split(':')
        assert len(tokens) == 3
        term = tokens[0].strip()
        atoms = [j.strip() for j in tokens[1].split(",")]
        key = ForceFieldKey(term, atoms)
        item = ForceFieldItem()
        item += [float(i) for i in v[2].split()]
        item.version = v[0]
        item.reference = v[3]
        return key, item
        
        
        
        


